/**
* @file post_process.cpp
*
* Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*/
#include "yolo/post_process.h"
#include <algorithm>
#include <bits/stdc++.h>
#include "string.h"
#include "unistd.h"

using namespace std;


namespace {
    // const static std::vector<std::string> kLabels = { "person", "bicycle", "car", "motorbike",
    //   "aeroplane","bus", "train", "truck", "boat",
    //   "traffic light", "fire hydrant", "stop sign", "parking meter",
    //   "bench", "bird", "cat", "dog", "horse",
    //   "sheep", "cow", "elephant", "bear", "zebra",
    //   "giraffe", "backpack", "umbrella", "handbag","tie",
    //   "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    //   "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    //   "tennis racket", "bottle", "wine glass", "cup",
    //   "fork", "knife", "spoon", "bowl", "banana",
    //   "apple", "sandwich", "orange", "broccoli", "carrot",
    //   "hot dog", "pizza", "donut", "cake", "chair",
    //   "sofa", "potted plant", "bed", "dining table", "toilet",
    //   "TV monitor", "laptop", "mouse", "remote", "keyboard",
    //   "cell phone", "microwave", "oven", "toaster", "sink",
    //   "refrigerator", "book", "clock", "vase","scissors",
    //   "teddy bear", "hair drier", "toothbrush" };

    const static std::vector<std::string> kLabels = { "car", "van", "bus", "others" };

    const std::vector<uint32_t> kClassToDetect = {0, 1, 2, 3};
    const size_t kClassNum = kClassToDetect.size();
    const size_t kOutputSize = kLabels.size() + 5; 

    const size_t kModelOutputBoxNum = 10647;
    const float kNMSThreshold = 0.5;
    const float kScoreThreshold = 0.7;


    // 是否在Mat上画框
    const bool kIsDraw = false;
    // 是否将检测结果存为图片落盘
    const bool kIsWrite = false;
    const char kStoragePath[] = "./output";

}

void PostProcess::SetBoxInfo(size_t index, BBox &box) {
    float *boxBuff = outputClass_;
    boxBuff += (index * kOutputSize);
    // box.rect.x = std::max<float>(std::min<float>(boxBuff[0], 1.0), 0.0) * width_;
    // box.rect.y = std::max<float>(std::min<float>(boxBuff[1], 1.0), 0.0) * height_;
    // box.rect.width = std::max<float>(std::min<float>(boxBuff[2], 1.0), 0.0) *width_;
    // box.rect.height = std::max<float>(std::min<float>(boxBuff[3], 1.0), 0.0) *height_;

    // center(x,y),w,h -> left(x,y),w,h
    float larger_width = boxBuff[2] * 1.1;
    float larger_height = boxBuff[3] * 1.1;
    box.rect.x = std::max<float>(std::min<float>((boxBuff[0]-larger_width/2.0), modelWidth_), 0);
    box.rect.y = std::max<float>(std::min<float>((boxBuff[1]-larger_height/2.0), modelHeight_), 0);
    box.rect.width = std::max<float>(std::min<float>(larger_width, modelHeight_ - box.rect.x), 0);
    box.rect.height = std::max<float>(std::min<float>(larger_height, imageHeight_ - box.rect.y), 0);
}


float PostProcess::IOU(const BBox &b1, const BBox &b2) {
    int x11 = std::max(b1.rect.x, b2.rect.x);
    int y11 = std::max(b1.rect.y, b2.rect.y);
    int x22 = std::min(b1.rect.x + b1.rect.width, b2.rect.x + b2.rect.width);
    int y22 = std::min(b1.rect.y + b1.rect.height, b2.rect.y + b2.rect.height);
    if (y11>y22 || x11>x22) {
        return 0.0f;
    }

    float sand = (x22 - x11) * (y22 - y11) * 1.0f;
    float area1 = b1.rect.width * b1.rect.height * 1.0f;
    float area2 = b2.rect.width * b2.rect.height * 1.0f;

    return sand / (area1 + area2 - sand);
}

void PostProcess::NMS(vector<BBox> &boxes, const float threshold) {

    std::sort(boxes.begin(), boxes.end(), 
                [](const BBox &a, const BBox &b) {return a.score < b.score;});

    for (size_t i = 0; i < boxes.size(); ++i) {
        boxes.erase(std::remove_if(boxes.begin() + i + 1, boxes.end(),
                                [&](const BBox &b) {
                                    return IOU(boxes[i], b) > threshold;
                                }),
                    boxes.end());
    }
    for (size_t i = 0; i < boxes.size(); ++i) {
        Utils::inv_letterbox_bbox(boxes[i].rect, cv::Size(modelWidth_, modelHeight_), cv::Size(imageWidth_, imageHeight_));   
    }
}


void PostProcess::DrawBoundBoxToImage_ycf(const vector<BBox> &result, cv::Mat& imageInput) {

    for (size_t i = 0; i < result.size(); ++i) {
        BBox box = result[i];
        string cls_name = kLabels[box.cls_id];
        cerr<<"the class name is: " + cls_name<<endl;
        fprintf(stderr, "%.2f %.2f %.2f %.2f %.2f\n", box.rect.x, box.rect.y, box.rect.width,box.rect.height,box.score);
        
        cv::Point origin;
        origin.x = box.rect.x;
        origin.y = box.rect.y;
        cv::Scalar fontColor = Utils::color_map(box.cls_id);
        cv::putText(imageInput, cls_name +  std::to_string(box.score) + "\%", origin, 0, 1, fontColor, 2, 4, 0);
        cv::rectangle(imageInput,  result[i].rect, fontColor, 1, 8, 0);
    }
    if (kIsWrite) {
        Utils::CheckFolder(kStoragePath);
        static int count = 0;
        string filename = string(kStoragePath) + "/result_" + to_string(count) + ".jpg";
        cv::imwrite(filename, imageInput);
        count++;
    }
}

void PostProcess::DrawDetections(cv::Mat& ori_img, vector<BBox>& detections) {
    for (auto &d : detections) {
        Utils::draw_bbox(ori_img, d.rect, kLabels[d.cls_id], Utils::color_map(d.cls_id+1));
    }
}

Result PostProcess::Process(cv::Mat& image, vector<BBox>& boxes) {

    imageHeight_ = image.rows;
    imageWidth_ = image.cols;

    xScale_ = (float)imageWidth_ / modelWidth_;
    yScale_ = xScale_;
    boxes.clear();
    float *classBuff = outputClass_;
    for (size_t i = 0; i < kModelOutputBoxNum; ++i) {
        float maxValue = 0;
        float maxIndex = 0;
        size_t j;
        for (j = 0; j < kClassNum; ++j) {
            uint32_t classId = kClassToDetect[j];
            float value = classBuff[i * kOutputSize + classId + 5];
            value *= classBuff[i * kOutputSize + 4]; // conf = obj_conf * cls_conf
            if (value > maxValue) {
                maxIndex = classId;
                maxValue = value;
            }
        }

        if (maxValue >= kScoreThreshold) {
            BBox b;
            SetBoxInfo(i, b);
            b.score = maxValue;
            b.cls_id = maxIndex;
            b.index = i;
            if (b.rect.width <= 0 || b.rect.height <= 0) 
                continue;
            boxes.push_back(b);
        }
    }

    
    NMS(boxes, kNMSThreshold);
    if (boxes.size()!=0 && kIsDraw) {
        DrawBoundBoxToImage_ycf(boxes, image);
    }
    return SUCCESS;
}
